# progress.py
"""进度显示模块，包含进度条和旋转动画功能"""

import sys
import time
import threading

class ProgressSpinner:
    """简单的进度旋转动画"""
    def __init__(self, message="处理中"):
        self.message = message
        self.spinner_chars = ["|", "/", "-", "\\"]
        self.stop_spinner = False
        self.thread = None
        
    def spin(self):
        """显示旋转动画"""
        i = 0
        while not self.stop_spinner:
            sys.stdout.write(f"\r{self.message} {self.spinner_chars[i % len(self.spinner_chars)]}")
            sys.stdout.flush()
            time.sleep(0.1)
            i += 1
        sys.stdout.write("\r" + " " * (len(self.message) + 2) + "\r")
        sys.stdout.flush()
    
    def start(self):
        """启动旋转动画"""
        self.stop_spinner = False
        self.thread = threading.Thread(target=self.spin)
        self.thread.daemon = True
        self.thread.start()
    
    def stop(self):
        """停止旋转动画"""
        self.stop_spinner = True
        if self.thread:
            self.thread.join()

class ProgressBar:
    """进度条显示"""
    def __init__(self, total_size, description="进度", min_interval=0.1):
        self.total_size = total_size
        self.description = description
        self.current_size = 0
        self.start_time = time.time()
        self.last_update_time = 0        
        self.min_interval = min_interval  # 最小更新间隔(秒)
        self.last_output_length = 0       # 记录上一次输出的长度

    def update(self, chunk_size):
        """更新进度"""
        self.current_size += chunk_size

        # 控制更新频率
        current_time = time.time()
        if current_time - self.last_update_time < self.min_interval and self.current_size < self.total_size:
            return
            
        self.last_update_time = current_time

        percent = (self.current_size / self.total_size) * 100 if self.total_size > 0 else 0
        
        # 计算速度
        elapsed_time = time.time() - self.start_time
        speed = self.current_size / elapsed_time if elapsed_time > 0 else 0
        
        # 计算剩余时间
        if speed > 0:
            remaining_time = (self.total_size - self.current_size) / speed
        else:
            remaining_time = 0
            
        # 格式化显示
        bar_length = 30
        filled_length = int(bar_length * percent / 100)
        bar = '█' * filled_length + '-' * (bar_length - filled_length)
        
        # 格式化大小和速度
        size_str = self._format_size(self.current_size)
        total_str = self._format_size(self.total_size)
        speed_str = self._format_size(speed) + "/s"

        # 构建输出字符串
        output = (
            f"\r{self.description}: |{bar}| {percent:.1f}% "
            f"[{size_str}/{total_str}] {speed_str} "
            f"ETA: {remaining_time:.1f}s"
        )

        # 清除上一次的输出（用空格覆盖）
        if len(output) < self.last_output_length:
            output += " " * (self.last_output_length - len(output))

        # 输出并刷新
        sys.stdout.write(output)
        sys.stdout.flush()
        
    def finish(self):
        """完成进度显示"""
        elapsed_time = time.time() - self.start_time
        speed = self.total_size / elapsed_time if elapsed_time > 0 else 0
        speed_str = self._format_size(speed) + "/s"

        # 完成后换行
        sys.stdout.write(
            f"\r{self.description}: 完成! "
            f"[{self._format_size(self.total_size)}] "
            f"用时: {elapsed_time:.1f}s ({speed_str})\n"
        )        
        # sys.stdout.write(
        #     f"\r{self.description}: 完成! "
        #     f"[{self._format_size(self.total_size)}] "
        #     f"用时: {elapsed_time:.1f}s ({speed_str})\n"
        # )
        sys.stdout.flush()
        
    def _format_size(self, size):
        """格式化文件大小"""
        for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
            if size < 1024.0:
                return f"{size:.3f} {unit}"
            size /= 1024.0
        return f"{size:.1f} PB"
